package comlib

import scala.collection.mutable.{ListBuffer}


class XTree[T](val node:T, val children:T=>Seq[T] ) {
  
  /**
   * 映射
   *
   * 按照leaf-root的方向遍历整个树，对于叶子节点执行actionl，
   * 对于非叶子节点执行actionb, 然后把返回值组成新的树
   *
   * @param actionb 非叶子节点执行动作。函数格式:
   *   - node: 当前遍历的节点
   *   - idx: 当前节点及父辈节点的序号
   *   - Nil: 子节点序列
   *   - stack: 当前节点和父辈节点序列
   *   - wides: 父辈节点兄弟个数列表
   * @param actionl 叶子节点执行动作
   *   - 当actionl=null时，actionl=actionb
   * @param nchildren 新树对应的children
   * @return 返回遍历执行完毕后生成的新树
   */
  def map[R](
    actionb : (T, Seq[Int], Seq[R], Seq[T], Seq[Int]) => R,
    actionl : (T, Seq[Int], Seq[R], Seq[T], Seq[Int]) => R = null,
    nchildren : R=>Seq[R] = null
  ):XTree[R] = {
    
    val actl = if( actionl==null ) actionb else actionl

    val res = XTree.map_iter( node, Nil, Nil, Nil, children, actionb, actl )
    new XTree( res, nchildren )
  }

  def isLeaf(n:T) = {
    val cs = children(n)
    return cs==Nil || cs==null
  }

  def find(criteria:(T,Seq[Int],Seq[Seq[T]],Seq[T],Seq[Int])=>Boolean) ={
    
    def recur(node:T, idx:Seq[Int], stack:Seq[T], wide:Seq[Int] ) : Seq[T] = {
      val childs = children(node)
      val sstack = stack :+ node

      if( childs!=null && childs != Nil ){
        val ret = childs.zipWithIndex map {case (c,i) => recur(c, idx:+i, sstack, wide:+childs.length)}
        if( criteria( node, idx, ret, sstack, wide ))
          sstack
        else
          ret.find(_!=Nil).getOrElse(Nil)
      }else{
        if( criteria( node, idx, Nil, sstack, wide )) sstack
        else Nil
      }
    }

    val res = recur( node, Nil, Nil, Nil )
    if( res == Nil ) None
    else Some(res)
  }

  /**
   * 遍历
   *
   * 按照direction指示的方向遍历整个树，对于叶子节点执行actionl，
   * 对于非叶子节点执行actionb
   *
   * @param actionb 非叶子节点执行动作。函数格式:
   *   - node: 当前遍历的节点
   *   - idx: 当前节点及父辈节点的序号
   *   - Nil: 子节点序列
   *   - stack: 当前节点和父辈节点序列
   *   - wides: 父辈节点兄弟个数列表
   * @param actionl 叶子节点执行动作
   *   - 当actionl=null时，actionl=actionb
   * @param direction 遍历方向
   *   - 0: Root-Leaf
   *   - 1: Leaf-Root
   * @return 返回遍历执行完毕后的this(XTree)
   */
  def foreach(
    actionb : (T, Seq[Int], Seq[T], Seq[T], Seq[Int]) => Unit,
    actionl : (T, Seq[Int], Seq[T], Seq[T], Seq[Int]) => Unit = null,
    direction : Int = 0
  ):XTree[T] = {
    
    val actl = if( actionl==null ) actionb else actionl

    def recur(
      node:T, 
      idx:Seq[Int], 
      stack:Seq[T], 
      wide:Seq[Int],
    ):Unit = {
      val childs = children(node)
      val sstack = stack :+ node

      if( childs != Nil && childs != null ){
        if( direction == 0 ) actionb( node, idx, childs, sstack, wide )
        childs.zipWithIndex foreach { case (c,i) =>
          recur( c, idx:+i, sstack, wide:+childs.length )
        }
        if( direction == 1 ) actionb( node, idx, childs, sstack, wide )
      }else{
        actl( node, idx, Nil, sstack, wide )
      }
    }

    recur( node, Nil, Nil, Nil )
    this
  }

  def foreach2(
    actionb : (T, Seq[Int], Seq[T], Seq[mutable.HashMap[String,Any]], Seq[Int]) => Unit,
    actionl : (T, Seq[Int], Seq[T], Seq[mutable.HashMap[String,Any]], Seq[Int]) => Unit = null,
    initStack : mutable.HashMap[String,Any] = mutable.HashMap(),
    direction : Int = 0
  ):XTree[T] = {
    
    val actl = if( actionl==null ) actionb else actionl

    def recur(
      node:T, 
      idx:Seq[Int], 
      stack:Seq[mutable.HashMap[String,Any]], 
      wide:Seq[Int],
    ):Unit = {
      val childs = children(node)
      val sstack = stack :+ mutable.HashMap("node"->node.asInstanceOf[Any])

      if( childs != Nil && childs != null ){
        if( direction == 0 ) actionb( node, idx, childs, sstack, wide )
        childs.zipWithIndex foreach { case (c,i) =>
          recur( c, idx:+i, sstack, wide:+childs.length )
        }
        if( direction == 1 ) actionb( node, idx, childs, sstack, wide )
      }else{
        actl( node, idx, Nil, sstack, wide )
      }
    }

    recur( node, Nil, Seq(initStack), Nil )
    this
  }

}

class MTree[T](val nodes:Seq[T], val children:T=>Seq[T] ) {

  /**
   * 映射
   *
   * 按照leaf-root的方向遍历整个树，对于叶子节点执行actionl，
   * 对于非叶子节点执行actionb, 然后把返回值组成新的树
   *
   * @param actionb 非叶子节点执行动作。函数格式:
   *   - node: 当前遍历的节点
   *   - idx: 当前节点及父辈节点的序号
   *   - Nil: 子节点序列
   *   - stack: 当前节点和父辈节点序列
   *   - wides: 父辈节点兄弟个数列表
   * @param actionl 叶子节点执行动作
   *   - 当actionl=null时，actionl=actionb
   * @param nchildren 新树对应的children
   * @return 返回遍历执行完毕后生成的新树
   */
  def map[R](
    actionb : (T, Seq[Int], Seq[R], Seq[T], Seq[Int]) => R,
    actionl : (T, Seq[Int], Seq[R], Seq[T], Seq[Int]) => R = null,
    nchildren : R=>Seq[R] = null
  ) = {
    val actl = if( actionl==null ) actionb else actionl
    val res = nodes.map( XTree.map_iter( _, Nil, Nil, Nil, children, actionb, actl ))
    new MTree( res, nchildren )
  }

  def isLeaf(n:T) = {
    val cs = children(n)
    return cs==Nil || cs==null
  }

  /**
   * 遍历
   *
   * 按照direction指示的方向遍历整个树，对于叶子节点执行actionl，
   * 对于非叶子节点执行actionb
   *
   * @param actionb 非叶子节点执行动作。函数格式:
   *   - node: 当前遍历的节点
   *   - idx: 当前节点及父辈节点的序号
   *   - Nil: 子节点序列
   *   - stack: 当前节点和父辈节点序列
   *   - wides: 父辈节点兄弟个数列表
   * @param actionl 叶子节点执行动作
   *   - 当actionl=null时，actionl=actionb
   * @param direction 遍历方向
   *   - 0: Root-Leaf
   *   - 1: Leaf-Root
   * @return 返回遍历执行完毕后的this(XTree)
   */
  def foreach(
    actionb : (T, Seq[Int], Seq[T], Seq[T], Seq[Int]) => Unit,
    actionl : (T, Seq[Int], Seq[T], Seq[T], Seq[Int]) => Unit = null,
    direction : Int = 0
  ):MTree[T] = {
    
    val actl = if( actionl==null ) actionb else actionl

    def recur(
      node:T, 
      idx:Seq[Int], 
      stack:Seq[T], 
      wide:Seq[Int],
    ):Unit = {
      val childs = children(node)
      val sstack = stack :+ node

      if( childs != Nil && childs != null ){
        if( direction == 0 ) actionb( node, idx, childs, sstack, wide )
        childs.zipWithIndex foreach { case (c,i) =>
          recur( c, idx:+i, sstack, wide:+childs.length )
        }
        if( direction == 1 ) actionb( node, idx, childs, sstack, wide )
      }else{
        actl( node, idx, Nil, sstack, wide )
      }
    }

    nodes.zipWithIndex.foreach {case (n,i) =>
      recur(n, Seq(i), Seq(initStack), Seq(nodes.length) )
    }
    this
  }
    
  def foreach2(
    actionb : (T, Seq[Int], Seq[T], Seq[mutable.HashMap[String,Any]], Seq[Int]) => Unit,
    actionl : (T, Seq[Int], Seq[T], Seq[mutable.HashMap[String,Any]], Seq[Int]) => Unit = null,
    initStack : mutable.HashMap[String,Any] = mutable.HashMap(),
    direction : Int = 0
  ):MTree[T] = {
    
    val actl = if( actionl==null ) actionb else actionl

    def recur(
      node:T, 
      idx:Seq[Int], 
      stack:Seq[mutable.HashMap[String,Any]], 
      wide:Seq[Int],
    ):Unit = {
      val childs = children(node)
      val sstack = stack :+ mutable.HashMap("node"->node.asInstanceOf[Any])

      if( childs != Nil && childs != null ){
        if( direction == 0 ) actionb( node, idx, childs, sstack, wide )
        childs.zipWithIndex foreach { case (c,i) =>
          recur( c, idx:+i, sstack, wide:+childs.length )
        }
        if( direction == 1 ) actionb( node, idx, childs, sstack, wide )
      }else{
        actl( node, idx, Nil, sstack, wide )
      }
    }

    nodes.zipWithIndex.foreach {case (n,i) =>
      recur(n, Seq(i), Seq(n), Seq(nodes.length) )
    }
    this
  }
    
}

object XTree{

  def apply[T](node:T, children:T=>Seq[T]) = new XTree[T](node, children)

  def recur[E,N](
    arr:Seq[E], 
    ofsti:Int, 
    result:ListBuffer[N],
    createNode:(E,Seq[N])=>N,
    level:E=>Int,
    currLvl:Int,
  ):Int = {
      
    var ofst = ofsti
    while( ofst < arr.length-1 ){
      val e = arr(ofst)
      val lvl = level(e)
      val nxt_lvl = level(arr(ofst+1))

      //if( lvl < currLvl ){
      //  return ofst
      //}else 
      if( lvl == nxt_lvl ){
        val ele = createNode(e,Nil)
        result.append(ele)
        ofst = ofst + 1
      }else if( lvl > nxt_lvl ){
        val ele = createNode(e, Nil)
        result.append(ele)
        return ofst + 1
      }else{
        var sub_result = ListBuffer[N]()
        ofst = recur( arr, ofst+1, sub_result, createNode, level, currLvl+1 )
        val ele = createNode( e, sub_result.toList )
        result.append( ele )

        if(ofst < arr.length-1){
          if( lvl > level(arr(ofst)) )
            return ofst
        }
      }
    }

    if( ofst == arr.length-1 ){
      val ele = createNode( arr(ofst), Nil )
      result.append( ele )
      ofst = ofst + 1
    }

    return ofst
  }


  /*========================================================================================================================
   构造函数：从迭代器中构造树结构
   ========================================================================================================================*/

  /** 从迭代器构建树结构
   * @tparam I 迭代器类型
   * @tparam E 输入迭代器的类型
   * @tparam N 输出“Tree”节点类型
   * @param ite 迭代器
   * @param level 获取节点在tree中的层次关系
   *  - e: {E} ---- 输入元素
   *  - return: {Int} ---- 层次
   * @param createNode 从迭代器构建树节点的函数
   *  - e: {E} ---- 输入迭代器
   *  - s: {Seq[N]} ---- children列表
   *  - return: {N} ---- 构造的当前树节点
   * @param initLvl 迭代器中Root元素对应的树层次
   * @return 树Root节点
   *
   * Note: 这里要求ite是Leaf-Root方向的遍历顺序
   */
  def fromIterator[I,E,N]( ite:I, level:E=>Int, createNode:(E,Seq[N])=>N, initLvl:Int=0 )( implicit fromIte:XIterable[I,E] ):N = MTree.fromIterator[I,E,N](ite, level, createNode, initLvl)(fromIte)(0)

  /*========================================================================================================================
   构造函数：从迭代器中构造树结构
   ========================================================================================================================*/
  def map_iter[N,R](
    node:N, idx:Seq[Int], stack:Seq[N], wide:Seq[Int],
    children : N=>Seq[N],
    actionb : (N,Seq[Int],Seq[R],Seq[N],Seq[Int]) => R,
    actionl : (N,Seq[Int],Seq[R],Seq[N],Seq[Int]) => R,
  ):R = {
    val childs = children(node)
    val sstack = stack :+ node

    if( childs != Nil && childs != null ){
      val ret = childs.zipWithIndex map { case (c,i) =>
        map_iter( c, idx:+i, sstack, wide:+childs.length, children, actionb, actionl )
      }
      actionb( node, idx, ret, sstack, wide )
    }else{
      actionl( node, idx, Nil, sstack, wide )
    }
  }

  def fromArray[E,N]( arr:Seq[E], createNode:(E,Seq[N])=>N, level:E=>Int, children:T=>Seq[T]=null)={
    var result = ListBuffer[N]()
    recur( arr, 0, result, createNode, level, 1 )
    new XTree(result(0), children)
  }

}

object MTree{
  def apply[T](node:T, children:T=>Seq[T]) = new MTree[T](node, children)

  def fromArray[E,N]( arr:Seq[E], createNode:(E,Seq[N])=>N, level:E=>Int, children:T=>Seq[T]=null) = {
    var result = ListBuffer[N]()
    XTree.recur( arr, 0, result, createNode, level, 1 )
    new MTree(result.toList, children)
  }

  /*========================================================================================================================
   构造函数：从迭代器中构造树结构
   ========================================================================================================================*/

  /** 从迭代器构建树结构
   * @tparam I 迭代器类型
   * @tparam E 输入迭代器的类型
   * @tparam N 输出“Tree”节点类型
   * @param ite 迭代器
   * @param level 获取节点在tree中的层次关系
   *  - e: {E} ---- 输入元素
   *  - return: {Int} ---- 层次
   * @param createNode 从迭代器构建树节点的函数
   *  - e: {E} ---- 输入迭代器
   *  - s: {Seq[N]} ---- children列表
   *  - return: {N} ---- 构造的当前树节点
   * @param initLvl 迭代器中Root元素对应的树层次
   * @return 树Root节点列表
   *
   * Note: 这里要求ite是Leaf-Root方向的遍历顺序
   */
  def fromIterator[I,E,N]( ite:I, level:E=>Int, createNode:(E,Seq[N])=>N, initLvl:Int=0 )( implicit fromIte:XIterable[I,E] ):Seq[N] = {

    /*
     * @param ite 迭代器
     * @param brothers 当前节点对应的兄弟节点列表
     * @param procLvl 当前处理函数对应的节点Level
     */
    def recur( 
      ite:I, 
      brothers:ListBuffer[N]=ListBuffer[N](), 
      procLvl:Int=0 
      ):Unit = {
      /* 以如下例子做算法思考：
        Root (lvl=0)
          |-- A (lvl=1)
          |   |-- A0 (lvl=2)
          |   |-- A1 (lvl=2)
          |   |-- A2 (lvl=2)
          |       |-- A2a (lvl=3)
          |       |-- A2b (lvl=3)
          |-- B (lvl=1)
              |-- B0 (lvl=2)
              |   |-- B0a (lvl=3)
              |   |-- B0b (lvl=3)
              |-- B1 (lvl=2)
       */

      var node = ite    // 循环节点变量定义

      while( fromIte.hasNext(node) ){   // 循环下一个节点
        val ele = fromIte.get(node)
        val lvl = level(ele)   // 获取节点层次
        val children = ListBuffer[N]()

        // 当前节点层次深于函数处理层次。例如：从Start -> B1; 从B1 -> B0b; 从B -> A2b
        if( lvl > procLvl ){
          // 构建当前节点父祖辈的兄弟列表和子列表
          recur( node, children, procLvl+1 )
        }

        // 当前节点层次浅于函数处理层次。例如：B0a->B0; B0->B; A2a->A2; A0->A; A->Root
        else if( lvl < procLvl ){
          // 返回上层节点。这时不处理返回后的节点，直到lvl==procLvl
          return
        }

        // 当前节点和处理节点匹配。例如：B0b->B0a; A2b->A2a; A2->A1; A1->A0
        // 构建当前节点，并加入到兄弟节点列表中
        val currNode = createNode( ele, children )
        brothers.append( currNode )
        
        node = fromIte.next(node)   // 迭代下一个元素
      }

    }

    val iite = fromIte.iter(ite)
    val brothers = ListBuffer[N]()
    recur( iite, brothers, initLvl )    // 必须初始化brothers和children，传递Nil回导致无法返回

    brothers    // Root节点藏在了Root节点brothers列表的第一个（也是唯一一个）
  }

}
